{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-06-17 10:16:23,395 - INFO - 输入的name:000_net, 输出的name:['082_convolutional', '094_convolutional', '106_convolutional']\n",
      "2020-06-17 10:16:23,398 - INFO - 开始Inference....\n"
     ]
    }
   ],
   "source": [
    "# coding: utf-8\n",
    "# author: hxy\n",
    "# 2019-12-10\n",
    "\n",
    "\"\"\"\n",
    "照片的inference；\n",
    "默认推理过程在CPU上；\n",
    "\"\"\"\n",
    "import os\n",
    "import time\n",
    "import logging\n",
    "import onnxruntime\n",
    "from darknet_api import process_img, get_boxes, draw_box\n",
    "\n",
    "\n",
    "# 定义日志格式\n",
    "def log_set():\n",
    "    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n",
    "\n",
    "\n",
    "# 加载onnx模型\n",
    "def load_model(onnx_model):\n",
    "    sess = onnxruntime.InferenceSession(onnx_model)\n",
    "    in_name = [input.name for input in sess.get_inputs()][0]\n",
    "    out_name = [output.name for output in sess.get_outputs()]\n",
    "    logging.info(\"输入的name:{}, 输出的name:{}\".format(in_name, out_name))\n",
    "\n",
    "    return sess, in_name, out_name\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    log_set()\n",
    "    input_shape = (416, 416) #修改为自己的\n",
    "    # anchors\n",
    "    anchors_yolo = [[(116, 90), (156, 198), (373, 326)], [(30, 61), (62, 45), (59, 119)],\n",
    "                    [(10, 13), (16, 30), (33, 23)]]\n",
    "    anchors_yolo_tiny = [[(81, 82), (135, 169), (344, 319)], [(10, 14), (23, 27), (37, 58)]]\n",
    "    session, inname, outname = load_model(onnx_model='yolov3_608.onnx')\n",
    "    logging.info(\"开始Inference....\")\n",
    "    # 照片的批量inference\n",
    "    img_files_path = 'samples/'\n",
    "    imgs = os.listdir(img_files_path)\n",
    "\n",
    "    logging.debug(imgs)\n",
    "    for img_name in imgs:\n",
    "        img_full_path = os.path.join(img_files_path, img_name)\n",
    "        logging.debug(img_full_path)\n",
    "        img, img_shape, testdata = process_img(img_path=img_full_path,\n",
    "                                               input_shape=input_shape)\n",
    "        s = time.time()\n",
    "        prediction = session.run(outname, {inname: testdata})\n",
    "\n",
    "        # logging.info(\"推理照片 %s 耗时：% .2fms\" % (img_name, ((time.time() - s)*1000)))\n",
    "        boxes = get_boxes(prediction=prediction,\n",
    "                          anchors=anchors_yolo,\n",
    "                          img_shape=input_shape)\n",
    "        draw_box(boxes=boxes,\n",
    "                 img=img,\n",
    "                 img_shape=img_shape)\n",
    "        logging.info(\"推理照片 %s 耗时：% .2fms\" % (img_name, ((time.time() - s)*1000)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: http://mirrors.aliyun.com/pypi/simple/\n",
      "Collecting onnxruntime\n",
      "  Downloading http://mirrors.aliyun.com/pypi/packages/5b/37/572986fb63e0df4e026c5f4c11f6a8977344293587b451d9210a429f5882/onnxruntime-1.3.0-cp36-cp36m-manylinux1_x86_64.whl (3.9 MB)\n",
      "\u001b[K     |████████████████████████████████| 3.9 MB 6.1 MB/s \n",
      "\u001b[?25hRequirement already satisfied: onnx>=1.2.3 in /home/zhongls/.conda/envs/pyzhongls/lib/python3.6/site-packages (from onnxruntime) (1.7.0)\n",
      "Requirement already satisfied: protobuf in /home/zhongls/.local/lib/python3.6/site-packages (from onnxruntime) (3.11.3)\n",
      "Requirement already satisfied: numpy>=1.16.6 in /home/zhongls/.local/lib/python3.6/site-packages (from onnxruntime) (1.18.3)\n",
      "Requirement already satisfied: typing-extensions>=3.6.2.1 in /home/zhongls/.conda/envs/pyzhongls/lib/python3.6/site-packages (from onnx>=1.2.3->onnxruntime) (3.7.4.2)\n",
      "Requirement already satisfied: six in /home/zhongls/.conda/envs/pyzhongls/lib/python3.6/site-packages (from onnx>=1.2.3->onnxruntime) (1.14.0)\n",
      "Requirement already satisfied: setuptools in /home/zhongls/.conda/envs/pyzhongls/lib/python3.6/site-packages (from protobuf->onnxruntime) (46.1.3.post20200330)\n",
      "Installing collected packages: onnxruntime\n",
      "Successfully installed onnxruntime-1.3.0\n"
     ]
    }
   ],
   "source": [
    "!pip install onnxruntime  -i http://mirrors.aliyun.com/pypi/simple/  --trusted-host mirrors.aliyun.com"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-pyzhongls]",
   "language": "python",
   "name": "conda-env-.conda-pyzhongls-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
